{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch_geometric\n",
    "from torch_geometric.datasets import Planetoid"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "use_cuda_if_available = False"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load the dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = Planetoid(root=\"tutorial1\",name= \"Cora\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Dataset properties"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(dataset)\n",
    "print(\"number of graphs:\\t\\t\",len(dataset))\n",
    "print(\"number of classes:\\t\\t\",dataset.num_classes)\n",
    "print(\"number of node features:\\t\",dataset.num_node_features)\n",
    "print(\"number of edge features:\\t\",dataset.num_edge_features)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Dataset shapes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(dataset.data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"edge_index:\\t\\t\",dataset.data.edge_index.shape)\n",
    "print(dataset.data.edge_index)\n",
    "print(\"\\n\")\n",
    "print(\"train_mask:\\t\\t\",dataset.data.train_mask.shape)\n",
    "print(dataset.data.train_mask)\n",
    "print(\"\\n\")\n",
    "print(\"x:\\t\\t\",dataset.data.x.shape)\n",
    "print(dataset.data.x)\n",
    "print(\"\\n\")\n",
    "print(\"y:\\t\\t\",dataset.data.y.shape)\n",
    "print(dataset.data.y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os.path as osp\n",
    "\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "from torch_geometric.nn import SAGEConv\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = dataset[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Net(torch.nn.Module):\n",
    "    def __init__(self):\n",
    "        super(Net, self).__init__()\n",
    "        \n",
    "        self.conv = SAGEConv(dataset.num_features,\n",
    "                             dataset.num_classes,\n",
    "                             aggr=\"max\") # max, mean, add ...)\n",
    "\n",
    "    def forward(self):\n",
    "        x = self.conv(data.x, data.edge_index)\n",
    "        return F.log_softmax(x, dim=1)\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "device = torch.device('cuda' if torch.cuda.is_available() and use_cuda_if_available else 'cpu')\n",
    "model, data = Net().to(device), data.to(device)\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train():\n",
    "    model.train()\n",
    "    optimizer.zero_grad()\n",
    "    F.nll_loss(model()[data.train_mask], data.y[data.train_mask]).backward()\n",
    "    optimizer.step()\n",
    "\n",
    "\n",
    "def test():\n",
    "    model.eval()\n",
    "    logits, accs = model(), []\n",
    "    for _, mask in data('train_mask', 'val_mask', 'test_mask'):\n",
    "        pred = logits[mask].max(1)[1]\n",
    "        acc = pred.eq(data.y[mask]).sum().item() / mask.sum().item()\n",
    "        accs.append(acc)\n",
    "    return accs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "best_val_acc = test_acc = 0\n",
    "for epoch in range(1,100):\n",
    "    train()\n",
    "    _, val_acc, tmp_test_acc = test()\n",
    "    if val_acc > best_val_acc:\n",
    "        best_val_acc = val_acc\n",
    "        test_acc = tmp_test_acc\n",
    "    log = 'Epoch: {:03d}, Val: {:.4f}, Test: {:.4f}'\n",
    "    \n",
    "    if epoch % 10 == 0:\n",
    "        print(log.format(epoch, best_val_acc, test_acc))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
